import torch


class BatchSequenceMemory(torch.nn.Module):
    def __init__(self, *, capacity: int, dim: int):
        super().__init__()
        self.capacity = capacity
        self.dim = dim

        self.size = torch.nn.Buffer()
        self.buffer = torch.nn.Buffer()

    def append(self, patch: torch.Tensor):
        dim = self.dim % patch.ndim

        patch_size = patch.shape[dim]
        if patch_size > self.capacity:
            patch = patch[*((slice(None),) * dim), -self.capacity:]
            patch_size = self.capacity

        if self.buffer.numel() == 0:
            batch_shape = patch.shape[:dim]
            value_shape = patch.shape[dim:][1:]
            buffer_shape = [*batch_shape, self.capacity, *value_shape]
            self.size = torch.zeros(batch_shape, dtype=torch.int64, device=patch.device)
            self.buffer = torch.empty(buffer_shape, dtype=patch.dtype, device=patch.device)
        else:
            self.buffer = self.buffer.roll(-patch_size, dims=dim)

        self.size += patch_size
        self.buffer[*((slice(None),) * dim), -patch_size:] = patch

    def reset(self, mask: torch.Tensor | None = None):
        if self.buffer.numel() == 0:
            return
        if mask is None:
            self.size.zero_()
            return
        self.size[mask] = 0

    def get(self) -> tuple[torch.Tensor, torch.Tensor] | tuple[None, None]:
        if self.buffer.numel() == 0:
            return None, None
        value = self.buffer
        indices = torch.arange(self.capacity, dtype=torch.int64, device=self.size.device)
        mask = (self.capacity - indices - 1) < self.size.unsqueeze(-1)
        return value, mask
